package site.ruolin.KeyGen;

import site.ruolin.KeyGen.Key.KeyPair;
import site.ruolin.KeyGen.Key.PrivateKey;
import site.ruolin.KeyGen.Key.PublicKey;
import site.ruolin.Util.MathUtil;

import java.math.BigInteger;

/**
 * 密钥生成
 * @author ruolin & 邱锋 on 2025/09/11
 * 用于产生加密数据的公钥（PK）和私钥（SK）
 */
public class KeyGen {
    private final int KeySize; // 密钥位数

    /**
     * 构造函数
     * @param keySize 密钥位数
     */
    public KeyGen(int keySize) {
        KeySize = keySize;
    }

    /**
     * 生成密钥对
     * 生成两个大素数 p 和 q ，并计算 n、g、λ、μ
     * @author ruolin
     * @return Paillier密钥对
     */
    public KeyPair generateKeyPair() {
        // 生成两个素数 p 和 q, 并检查是否符合要求
        BigInteger[] pq = checkPrime(KeySize);
        BigInteger p = pq[0];
        BigInteger q = pq[1];

        // 计算 n = p * q
        BigInteger n = p.multiply(q);

        // 计算 g
        BigInteger g = n.add(BigInteger.ONE);

        // 计算 λ = lcm(p - 1, q - 1)
        BigInteger lambda = MathUtil.lcm(p.subtract(BigInteger.ONE), q.subtract(BigInteger.ONE));

        // 计算 μ = (L(g^λ mod n²))^(-1) mod n
        BigInteger nSquare = n.multiply(n);
        BigInteger gLambda = g.modPow(lambda, nSquare);
        BigInteger L_result = MathUtil.L(gLambda, n);
        BigInteger mu = L_result.modInverse(n);

        // 创建公钥和私钥
        PublicKey publicKey = new PublicKey(n, g);
        PrivateKey privateKey = new PrivateKey(lambda, mu);

        return new KeyPair(publicKey, privateKey);
    }

    /**
     * 检查生成的素数（p、q）是否符合满足 gcd(pq,(p - 1)(q - 1))=1，且满足 p 和 q 的长度相等;
     * @author 邱锋
     * @param KeySize 密钥位数
     * @return [0]:p [1]:q
     */
    private BigInteger[] checkPrime(int KeySize) {
        while (true) {
            BigInteger p = MathUtil.GeneratorPrime(KeySize);
            BigInteger q = MathUtil.GeneratorPrime(KeySize);

            // gcd(pq,(p - 1)(q - 1))
            BigInteger pq = p.multiply(q); // pq
            BigInteger p_1 = p.subtract(BigInteger.ONE); // p - 1
            BigInteger q_1 = q.subtract(BigInteger.ONE); // q - 1

            // gcd(pq,(p - 1)(q - 1)) == 1
            if (pq.gcd(p_1.multiply(q_1)).equals(BigInteger.ONE)) {
                return new BigInteger[]{p, q};
            }
        }
    }
}
